package ai.mspbots.poc.db.util;

import ai.mspbots.poc.db.entity.IndArrayVector;
import ai.mspbots.poc.db.service.IndArrayVectorService;
import ai.mspbots.poc.help.PocEnv;
import cn.hutool.extra.spring.SpringUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;

import java.util.Arrays;

public class INDArrayUtil {
    private static final IndArrayVectorService indArrayVectorService = SpringUtil.getBean(IndArrayVectorService.class);

    public static void saveToDb(String text, INDArray array) {

        if (array == null || array.isEmpty()) {
            return;
        }

        indArrayVectorService.save(new IndArrayVector(text,
                Arrays.stream(array.toDoubleVector()).boxed().toArray(Double[]::new)));
        //
        // // 转换为 Object[]
        // Object[] arrayData = new Object[(int) array.length()];
        // for (int i = 0; i < array.length(); i++) {
        //     arrayData[i] = array.getDouble(i); // 获取每个元素
        // }
        //
        // String sql = "INSERT INTO ind_array_vector(id, data) VALUES (?, ?)";
        // jdbcTemplate.update(sql, ps -> {
        //     ps.setString(1, text);
        //     ps.setArray(2, ps.getConnection().createArrayOf("float8", arrayData));
        // });
    }

    public static INDArray retrieveINDArray(String key) {

        IndArrayVector indArrayVector = indArrayVectorService.getById(key);
        if (indArrayVector != null) {
            Double[] data = indArrayVector.getData();
            if (data != null && data.length > 0) {
                return Nd4j.create(Arrays.asList(data));
            }
        }
        return Nd4j.zeros(SpringUtil.getBean(PocEnv.class).getDims());
    }

    public static int checkSql(String id) {
        return (int) indArrayVectorService.count(Wrappers.lambdaQuery(IndArrayVector.class).eq(IndArrayVector::getId, id));
        // String sql = "SELECT count(id) FROM ind_array_vector WHERE id = '" + id + "'";
        // return jdbcTemplate.query(sql, rs -> {
        //     if (!rs.next()) {
        //         return 0;
        //     }
        //     return rs.getInt(1);
        // });
    }

}
